from diffusers import AutoPipelineForText2Image, StableDiffusionPipeline
import torch
from PIL import Image

# # 初始化管道
# pipe = AutoPipelineForText2Image.from_pretrained(
#     "stabilityai/sd-turbo",
#     torch_dtype=torch.float16,
#     variant="fp16"
# )

# 初始化管道
pipe = StableDiffusionPipeline.from_single_file(
    "D:\\ComfyUI_windows_portable_nvidia\\ComfyUI\\models\\checkpoints\\1_5_model.safetensors",
    torch_dtype=torch.float16,
    # 移除 variant 参数
    # config 参数应指向包含模型配置的目录（非文件路径）
    config="D:\\SDServer\\stable-diffusion-server\\stabilityai\\sd-turbo"
)
pipe.to("cuda")

def generate_image_from_parameters(prompt, width=None, height=None, **kwargs):
    print('接受请求。')
    img = pipe(prompt=prompt, width=width, height=height, **kwargs)[0][0]
    # pipe.to("cpu")
    torch.cuda.empty_cache()
    return img
    # return Image.open("generated_image.jpg")